/*
 *  This file is part of Bracket Security
 *  Copyright 2011 David R. Smith
 *
 */
package asia.redact.bracket.security;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.nio.charset.Charset;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.Provider;
import java.security.PublicKey;
import java.security.Security;
import java.security.Signature;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateCrtKey;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.apache.log4j.Logger;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import asia.redact.bracket.security.config.SecurityConfig;
import asia.redact.bracket.security.config.ZeroConfig;
import asia.redact.bracket.util.SafeString;

/**
 * <pre>
 * Utility for doing digital signature, encryption, etc. Basically just a wrapper around
 * some BouncyCastle functionality.
 * 
 * What we're implementing is the quintessential algorithm which uses two
 * ciphers, one symmetric and the other asymmetric. We encrypt the message with the
 * symmetric cipher, then we encrypt the key for that cipher with RSA.  Both the key and 
 * the message is then armored with a URL-SAFE Base64 variant. 
 * 
 * For this class to work as expected, install full strength cryptographic Jurisdiction files 
 * into your JDK/JRE.
 * 
 * This class is thread-safe. It can be used as a static Singleton, etc.
 * 
 * This class fails with RuntimeExceptions and all the crypto checked exceptions are wrapped 
 * in this way. (If something is wrong with the crypto config, we assume it will not be possible to
 * fix at runtime.) 
 * 
 * </pre>
 * @author Dave
 *
 */
public final class SecurityTool {
	
	private static final Logger log = Logger.getLogger(SecurityTool.class);
	
	private final SecurityConfig config;
	private KeyStore ks; 
	private SecretKeySpec keySpec;
	private Cipher symmetricCipher;
	private Cipher asymmetricCipher;
	private final Lock lock = new ReentrantLock();
	
	/**
	 * Use the ZeroConfig module
	 */
	public SecurityTool() {
		lock.lock();
		try {
			config = new ZeroConfig();
			init();
		}finally{
			lock.unlock();
		}
	}
	
	public SecurityTool(SecurityConfig config) {
		lock.lock();
		try {
			this.config = config;
			if(config.useCaching()){
				initCachedMode();
			}else{
				init();
			}
		}finally{
			lock.unlock();
		}
	}
	
	private void init() {
		lock.lock();
		try {
			initProvider();
			initKeystore();
			initCiphers();
		}finally{
			lock.unlock();
		}
	}
	
	/**
	 * Not doing anything different at the moment - TODO implement cache mode
	 */
	private void initCachedMode() {
		lock.lock();
		try {
			initProvider();
			initKeystore();
			initCiphers();
		}finally{
			lock.unlock();
		}
	}
	

	private void initKeystore() {
		lock.lock();
		try {
		
			FileInputStream fis = null;
			try {
				ks = KeyStore.getInstance(getKeystoreType());
			    fis = new FileInputStream(getKeystoreFile());
			    ks.load(fis, getStorepassPassword());
			} catch (Exception e) {
				log.error("Keystore problem", e);
				throw new RuntimeException("Bailing out, something wrong with keystore or password...");
			}finally{
				if(fis != null)
					try {
						fis.close();
					} catch (IOException e) {
						log.error("IOException, generated keystore may not be usable", e);
					}	
			}
		}finally{
			lock.unlock();
		}
		
		log.info("*************Initialized keystore OK***********");
	}
	
	private void initCiphers() {
		lock.lock();
		try {
			try {
					keySpec = generateSymmetricKey(getSymmetricAlgorithm());
					log.info("initialized keySpec for "+getSymmetricAlgorithm());
			} catch (Exception e) {
					log.error("Failed to init", e);
					throw new RuntimeException("Bailing out, cannot generate a symmetric key!");
			}
		   
			try {
				symmetricCipher = Cipher.getInstance(getSymmetricAlgorithm(), getJCEProvider());
				log.info("initialized symmetric Cipher for "+getSymmetricAlgorithm()+" "+getJCEProvider());
			} catch (NoSuchProviderException e) {
				log.error("No such provider! You need Bouncy Castle!", e);
				throw new RuntimeException("Bailing out");
			} catch (NoSuchPaddingException e) {
				log.error("Padding error, you need full strength crypto", e);
				throw new RuntimeException("Bailing out");
			} catch (NoSuchAlgorithmException e) {
				log.error("UnknownAlgorithm", e);
				throw new RuntimeException("Bailing out");
			}
			   
			try {
				asymmetricCipher = Cipher.getInstance(getAsymmetricAlgorithm(), getJCEProvider());
				log.info("initialized symmetric Cipher for "+getAsymmetricAlgorithm()+" "+getJCEProvider());
			} catch (NoSuchProviderException e) {
				log.error("No such provider! Please install Bouncy Castle!", e);
				throw new RuntimeException("Bailing out");
			} catch (NoSuchPaddingException e) {
				log.error("Padding error, you need full strength crypto", e);
				throw new RuntimeException("Bailing out");
			} catch (NoSuchAlgorithmException e) {
				log.error("UnknownAlgorithm", e);
				throw new RuntimeException("Bailing out");
			}
			   
		}finally{
			lock.unlock();
		}
		
		log.info("************Ciphers initialized**************");
	}

	/**
	 * Add the BouncyCastle Provider dynamically if it is not found.
	 * 
	 */
	private void initProvider(){
		Provider [] providers = Security.getProviders();
		for(Provider p: providers){
			if(p instanceof BouncyCastleProvider){
				log.info("Found the Bouncy Castle provider in the config");
				return;
			}
		}
		log.warn("Adding the BouncyCastle provider dynamically. You might want to install it statically into the java security configuration for better performance.");
		Security.addProvider(new BouncyCastleProvider());
	}
	
	/**
	 * Check if the alias exists in the configuration keystore
	 * @return
	 */
	
	public boolean aliasExists(String token) {
		lock.lock();
		try {
			Enumeration<String> aliases;
			
			try {
				aliases = ks.aliases();
				while(aliases.hasMoreElements()){
					String alias = aliases.nextElement();
					log.debug("Next Alias: "+alias);
					if(alias.equals(token)){
						log.debug("Found alias: "+token);
						return true;
					}
				}
			} catch (KeyStoreException e) {
				log.error("Issue with keystore, could not find alias on those grounds", e);
			}
			log.debug("Alias "+token+" not found");
			return false;
		}finally{
			lock.unlock();
		}
	}
	
	/**
	 * Find the alias in the keystore for a given DN. It works better to pass in a partial DN as the value returned
	 * by Bouncy Castle's implementation has no guaranteed order as to the DN's components
	 * 
	 * Returns null on failure
	 * 
	 * @param distinguishedName
	 * @return
	 */
	public String findAlias(String distinguishedName) {
		lock.lock();
		try {
			try {
				Enumeration<String> aliases = ks.aliases();
				while(aliases.hasMoreElements()){
					String alias = aliases.nextElement();
					Certificate cert = ks.getCertificate(alias);
					X509Certificate x509 = (X509Certificate) cert;
					log.debug("looking at x509 principal: "+x509.getIssuerX500Principal().getName());
					if(x509.getIssuerX500Principal().getName().contains(distinguishedName)){
						log.info("Found '"+distinguishedName+"' within the principal of the x509 cert");
						return alias;
					}
				}
			} catch (KeyStoreException e) {
				log.error("Issue with Keystore, could not find alias on those grounds", e);
				throw new RuntimeException(e);
			}
			log.warn("Cound not find the DN '"+distinguishedName+"' - returning null");
		    return null;
		    
		}finally{
			lock.unlock();
		}
	}
	
	/**
	 * Find the alias in the keystore for a given certificate serial number
	 * Can return null on failure
	 * 
	 * @param distinguishedName
	 * @return
	 */
	public String findAlias(BigInteger serialNumber) {
		lock.lock();
		try {
			try {
				Enumeration<String> aliases = ks.aliases();
				while(aliases.hasMoreElements()){
					String alias = aliases.nextElement();
					log.debug("Checking "+alias+" for serial number "+serialNumber.toString());
					Certificate cert = ks.getCertificate(alias);
					X509Certificate x509 = (X509Certificate) cert;
					if(x509.getSerialNumber().equals(serialNumber)){
						log.info("found it, returning "+alias);
						return alias;
					}
				}
			} catch (KeyStoreException e) {
				log.error("Issue with keystore, could not find alias on those grounds", e);
			}
			log.warn("Did not find serialNumber "+serialNumber.toString()+" - returning null");
			return null;
		}finally{
			lock.unlock();
		}
	}
	
	/**
	 * TODO - given a file containing an X509 certificate, verify a message
	 * 
	 * @param x509CertFile
	 * @param message
	 * @return
	 */
	public boolean verify(File x509CertFile, SafeString signature, String message){
		return false;
	}
	
	/**
	 * TODO - given a file containing an X509 certificate, encrypt a message
	 * 
	 * @param x509CertFile
	 * @param message
	 * @return
	 */
	public EncryptResult encrypt(File x509CertFile, String message){
		return null;
	}
	
	
	
	/**
	 * Encrypt a message and return an armored symmetric key and message
	 * 
	 * @param params
	 * @return
	 * @throws Exception
	 */
	public EncryptResult encrypt(EncryptParams params) throws Exception{
		lock.lock();
		try {
			EncryptResult result = null;
			SecretKeySpec spec = null;
			if(!params.hasKey()){
				log.info("encrypting without passed in key, generating one");
				spec = generateSymmetricKey(params.symmetricAlgorithm);
				params.symmetricKey = spec;
				log.info("Generated: "+spec);
			}
			
				log.info("encrypting with key "+params.symmetricKey+" of size "+spec.getEncoded().length);
				result = encryptSymmetric(params);
				log.info("***************Complete with encrypt()*************");
			    return result;
		}finally{
			lock.unlock();
		}
	}
	
	public String decrypt(DecryptParams params) throws Exception{
		lock.lock();
		try {
			log.info("***********Start decrypt **************");
			log.info("params "+params.toString());
			
			byte [] key = decryptKeyAsymmetric(params.alias, params.encryptedKey);
			return decryptMessage(key, params.encryptedMsg, this.getSymmetricAlgorithm());
			
		}finally{
			lock.unlock();
		}
	}
	
	/**
	 * Returns a base64 url-safe encoded string which is the armored hash of the message encrypted using RSA.
	 * 
	 * TODO This code assumes an RSA signature try to generalize
	 * 
	 * @param alias
	 * @param message
	 * @return
	 */
	public SafeString sign(String alias, String message) {
		lock.lock();
		try {
		    if(ks.containsAlias(alias)) {
		    	Key key = ks.getKey(alias, getStorepassPassword());
		    	RSAPrivateCrtKey _key = (RSAPrivateCrtKey) key;
		    	Signature signature = Signature.getInstance(this.getSignatureAlgorithm(), this.getJCEProvider());
		    	// generate a signature
		        signature.initSign(_key);
		        byte[] bytes = message.getBytes();
		        signature.update(bytes);
		        byte[]  sigBytes = signature.sign();
		        
		       return new SafeString(sigBytes);
		        
		    }else{
		    	throw new RuntimeException("No alias "+alias+" found, you need to fix this to proceed");
		    }
		} catch (Exception e) {
			throw new RuntimeException("Failed to sign", e);
		}finally{
			lock.unlock();
		}
		
	}

	public boolean verify(String alias, String message, SafeString sig) {
		lock.lock();
		try {
		    if(ks.containsAlias(alias)) {
		    	Certificate cert = ks.getCertificate(alias);
		    //	Key key = ks.getKey(alias, getPassword());
		    	PublicKey pub = cert.getPublicKey();
		    	Signature signature = Signature.getInstance(this.getSignatureAlgorithm(), this.getJCEProvider());
		  
		        // verify a signature
		        signature.initVerify(pub);
		      //  byte [] sigBytes = Base64.decode(sig, Base64.NO_OPTIONS);
		        signature.update(message.getBytes());
		      
		        if (signature.verify(sig.decode())){
		        	log.info("Signature was verified - OK");
		            return true;
		        }else{
		        	log.info("Signature failed verification!");
		           return false;
		        }
		    }else{
		    	throw new RuntimeException("No alias "+alias+" found, you need to fix this to proceed");
		    }
		} catch (Exception e) {
			throw new RuntimeException("Failed to verify either way - error", e);
		}finally{
			lock.unlock();
		}
	}

	/**
	 * TODO Make this a defensive copy of the security config
	 * 
	 * @return
	 */
	public SecurityConfig getSecurityConfig() {
		return config;
	}

	private SecretKeySpec generateSymmetricKey(String keyType) throws Exception{
		
		lock.lock();
		try {
		   String provider = this.getJCEProvider();
	       KeyGenerator kgen = KeyGenerator.getInstance(keyType, provider);
	       kgen.init(this.getSymmetricKeySize());
	       SecretKey skey = kgen.generateKey();
	       byte[] raw = skey.getEncoded();
	       log.debug("generated symmetric key of type "+keyType);
	       return new SecretKeySpec(raw, keyType);
		}finally{
			lock.unlock();
		}
	}
	
	private EncryptResult encryptSymmetric(EncryptParams params) throws Exception {
		
		lock.lock();
		
		try {
			
			log.debug("Starting encryptSymmetric using "+params);
			
			if(symmetricCipher == null) {
				log.debug("symmetricCipher is null, getting one");
				symmetricCipher = Cipher.getInstance(params.symmetricAlgorithm, this.getJCEProvider());
			}
		    symmetricCipher.init(Cipher.ENCRYPT_MODE, params.symmetricKey);
		    byte[] encrypted = symmetricCipher.doFinal(params.message.getBytes());
		    log.info("encrypted "+encrypted.length+" bytes");
		  
		    // the key, Base64 encoded
		    SafeString armoredKey = encryptKeyAsymmetric(params.alias, params.symmetricKey.getEncoded());
		    EncryptResult result = new EncryptResult(armoredKey,new SafeString(encrypted));
		    
		    log.debug("done with encrypt, result: "+result);
		    
			return result;
			
		}finally{
			lock.unlock();
		}
	}
	
	private String decryptMessage(final byte [] key, final SafeString message, final String keyType) throws Exception {
		
		lock.lock();
		try {
			log.debug("key: "+key.length);
			log.debug("message: "+message);
			log.debug("keyType: "+keyType);
			SecretKeySpec spec =  new SecretKeySpec(key, keyType);
			if(symmetricCipher==null) {
				symmetricCipher = Cipher.getInstance(keyType, this.getJCEProvider());
			}
		    symmetricCipher.init(Cipher.DECRYPT_MODE, spec);
		    byte[] original = symmetricCipher.doFinal(message.decode());
		    String txt = new String(original, Charset.forName("UTF-8"));
		    log.debug("done with decrypt, returning message of length "+txt.length());
		    return txt;
		}finally{
			lock.unlock();
		}
	}
	
	private SafeString encryptKeyAsymmetric(String alias, byte [] keyBytesToEncrypt) throws Exception{
		
		lock.lock();
		try {
			
			log.info("Starting encryptAsymmetric of the cipher key");
			
			 	Certificate cert = ks.getCertificate(alias);
				if(asymmetricCipher == null) {
					asymmetricCipher = Cipher.getInstance(this.getAsymmetricAlgorithm(), this.getJCEProvider());
				}
		    	asymmetricCipher.init(Cipher.ENCRYPT_MODE,cert);
		    	byte [] rsaEnc = asymmetricCipher.doFinal(keyBytesToEncrypt);
		    	log.debug("Done with asymmetric encrypt on key, base64 encoding now...");
		    	return new SafeString(rsaEnc);
			
		}finally{
			lock.unlock();
		}
		
	}
	
	private byte [] decryptKeyAsymmetric(String alias, SafeString encodedKeyToDecrypt) throws Exception{
			lock.lock();
			try {
				log.info("Starting decrypt of symmetric key using "+alias+" for the RSA key");
				byte [] keyBytes = encodedKeyToDecrypt.decode();
				log.debug("keyBytes size: "+keyBytes.length);
				
			 	PrivateKey key = (PrivateKey) ks.getKey(alias, getKeyPassword());
			 	log.debug("private key loaded using alias: "+alias);
				if(asymmetricCipher == null) {
					asymmetricCipher = Cipher.getInstance(this.getAsymmetricAlgorithm(), this.getJCEProvider());
					log.debug("creating asymmetric cipher: "+asymmetricCipher.getAlgorithm());
				}
				log.debug("initializing cipher in decrypt mode using key "+key.getAlgorithm()+" "+key.getFormat());
				asymmetricCipher.init(Cipher.DECRYPT_MODE,key);
				log.debug("doing doFinal()");
		    	byte [] bytes = asymmetricCipher.doFinal(keyBytes);
		    	log.info("Done with RSA decrypt on symmetric key, returning raw bytes...");
		    	return bytes;
				
				
			}finally{
				lock.unlock();
			}
		}
	
	private char [] getStorepassPassword() {
		return config.getStorepass().toCharArray();
	}
	
	private char [] getKeyPassword() {
		return config.getKeypass().toCharArray();
	}
	
	private String getKeystoreType() {
		return config.getKeystoreType();
	}
	
	private String getSignatureAlgorithm() {
		return config.getSignatureAlgorithm();
	}
	
	private String getJCEProvider() {
		return config.getJCEProvider();
	}
	
	private String getSymmetricAlgorithm() {
		return config.getSymmetricAlgorithm();
	}
	
	private int getSymmetricKeySize() {
		return config.getSymmetricKeySize();
	}
	
	private String getAsymmetricAlgorithm() {
		return config.getAsymmetricAlgorithm();
	}
	
	/**
	 * Assumes bracket keystore is named "bracketKeystore" but you can set the extension for different types of keystore
	 * 
	 * @return
	 */
	private File getKeystoreFile() {
		lock.lock();
		try {
			File file = new File(config.getKeyMaterialsBase(), "bracketKeystore."+config.getKeystoreFileExtension());
			if(!file.exists()) {
				try {
					throw new RuntimeException("Keystore file "+file.getCanonicalPath()+" does not exist!");
				} catch (IOException e) {
					log.error("Did not find keystore, this is a problem you need to fix to proceed", e);
				}
			}
			return file;
		}finally{
			lock.unlock();
		}
		
	}
	
	@Override
	public int hashCode() {
		final int prime = 31;
		int result = 1;
		result = prime
				* result
				+ ((asymmetricCipher == null) ? 0 : asymmetricCipher.hashCode());
		result = prime * result + ((config == null) ? 0 : config.hashCode());
		result = prime * result + ((keySpec == null) ? 0 : keySpec.hashCode());
		result = prime * result + ((ks == null) ? 0 : ks.hashCode());
		result = prime * result + ((lock == null) ? 0 : lock.hashCode());
		result = prime * result
				+ ((symmetricCipher == null) ? 0 : symmetricCipher.hashCode());
		return result;
	}

	@Override
	public boolean equals(Object obj) {
		if (this == obj) {
			return true;
		}
		if (obj == null) {
			return false;
		}
		if (!(obj instanceof SecurityTool)) {
			return false;
		}
		SecurityTool other = (SecurityTool) obj;
		if (asymmetricCipher == null) {
			if (other.asymmetricCipher != null) {
				return false;
			}
		} else if (!asymmetricCipher.equals(other.asymmetricCipher)) {
			return false;
		}
		if (config == null) {
			if (other.config != null) {
				return false;
			}
		} else if (!config.equals(other.config)) {
			return false;
		}
		if (keySpec == null) {
			if (other.keySpec != null) {
				return false;
			}
		} else if (!keySpec.equals(other.keySpec)) {
			return false;
		}
		if (ks == null) {
			if (other.ks != null) {
				return false;
			}
		} else if (!ks.equals(other.ks)) {
			return false;
		}
		if (lock == null) {
			if (other.lock != null) {
				return false;
			}
		} else if (!lock.equals(other.lock)) {
			return false;
		}
		if (symmetricCipher == null) {
			if (other.symmetricCipher != null) {
				return false;
			}
		} else if (!symmetricCipher.equals(other.symmetricCipher)) {
			return false;
		}
		return true;
	}

}
